# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Common functions
"""

from typing import Union, List, Tuple, Iterable
import numpy as np
from numpy import ndarray
import torch
from torch import Tensor, nn
from torch.nn import Parameter
import h5py
from joblib import Parallel, delayed

__all__ = [
    'PI',
    'keepdims_sum',
    'keepdims_mean',
    'keepdims_prod',
    'reduce_any',
    'reduce_all',
    'reduce_prod',
    'concat_first_dim',
    'concat_last_dim',
    'concat_penulti',
    'stack_first_dim',
    'stack_last_dim',
    'stack_penulti',
    'squeeze_first_dim',
    'squeeze_last_dim',
    'squeeze_penulti',
    'periodic_variable',
    'periodic_difference',
    'gather_vector',
    'gather_value',
    'pbc_box_reshape',
    'pbc_image',
    'coordinate_in_pbc',
    'vector_in_pbc',
    'calc_vector_nopbc',
    'calc_vector_pbc',
    'calc_vector',
    'calc_distance_nopbc',
    'calc_distance_pbc',
    'calc_distance',
    'calc_angle_by_vectors',
    'calc_angle_nopbc',
    'calc_angle_pbc',
    'calc_angle',
    'calc_torsion_by_vectors',
    'calc_torsion_nopbc',
    'calc_torsion_pbc',
    'calc_torsion',
    'coulomb_interaction',
    'lennard_jones_potential',
    'lennard_jones_potential2',
    'get_integer',
    'get_ndarray',
    'get_tensor',
    'check_broadcast',
    'any_none',
    'all_none',
    'any_not_none',
    'all_not_none',
    'get_arguments',
    'bonds_in',
    'calc_gaussian',
    'H5MDProcessor',
    'get_global_device',
    'set_torchsponge_global_device',
]

PI = 3.141592653589793238462643383279502884197169399375105820974944592307
r""":math:`\pi`"""

def keepdims_sum(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:
    """
    Reduces a dimension to 1 by summing the elements in the dimension of `x` along the axis,
    and the dimensions of the output and input are the same.

    Args:
         - **x** (Tensor[Number]) - The input tensor. The dtype of the tensor to be reduced is number.
           :math:`(N,*)` where :math:`*` means, any number of additional dimensions, its rank should be less than 8.
         - **axis** (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.
           Only constant value is allowed. Must be in the range [-rank(`x`), rank(`x`)).

    Outputs:
        Tensor, has the same dtype as the `x`.
    """
    if isinstance(axis, int):
        axis = (axis,)
    return torch.sum(x, dim=axis, keepdim=True)

def keepdims_mean(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:
    """
    Reduces a dimension to 1 by averaging the elements in the dimension of `x` along the axis,
    and the dimensions of the output and input are the same.

    Args:
         - **x** (Tensor[Number]) - The input tensor. The dtype of the tensor to be reduced is number.
           :math:`(N,*)` where :math:`*` means, any number of additional dimensions, its rank should be less than 8.
         - **axis** (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.
           Only constant value is allowed. Must be in the range [-rank(`x`), rank(`x`)).

    Outputs:
        Tensor, has the same dtype as the `x`.
    """
    if isinstance(axis, int):
        axis = (axis,)
    return torch.mean(x, dim=axis, keepdim=True)

def keepdims_prod(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:
    """
    Reduces a dimension to 1 by multiplying the elements in the dimension of `x` along the axis,
    and the dimensions of the output and input are the same.

    Args:
         - **x** (Tensor[Number]) - The input tensor. The dtype of the tensor to be reduced is number.
           :math:`(N,*)` where :math:`*` means, any number of additional dimensions, its rank should be less than 8.
         - **axis** (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.
           Only constant value is allowed. Must be in the range [-rank(`x`), rank(`x`)).

    Outputs:
        Tensor, has the same dtype as the `x`.
    """
    if isinstance(axis, int):
        axis = (axis,)
    return torch.prod(x, dim=axis, keepdim=True)

def reduce_any(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:
    r"""
    Reduces a dimension of a tensor by the "logical OR" of all elements in the dimension, by default. And also can
    reduce a dimension of `x` along the axis.

    Args:
        - **x** (Tensor[bool]) - The input tensor. The dtype of the tensor to be reduced is bool.
          :math:`(N,*)` where :math:`*` means, any number of additional dimensions, its rank should be less than 8.
        - **axis** (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.
          Only constant value is allowed. Must be in the range [-rank(x), rank(x)).

    Outputs:
        Tensor, the dtype is bool.
    """
    if isinstance(axis, int):
        axis = (axis,)
    return torch.any(x, dim=axis)

def reduce_all(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:
    r"""
    Reduces a dimension of a tensor by the "logical AND" of all elements in the dimension, by default. And also can
    reduce a dimension of `x` along the axis.

    Args:
        - **x** (Tensor[bool]) - The input tensor. The dtype of the tensor to be reduced is bool.
          :math:`(N,*)` where :math:`*` means, any number of additional dimensions, its rank should be less than 8.
        - **axis** (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.
          Only constant value is allowed. Must be in the range [-rank(x), rank(x)).

    Outputs:
        Tensor, the dtype is bool.
    """
    if isinstance(axis, int):
        axis = (axis,)
    return torch.all(x, dim=axis)

def reduce_prod(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:
    r"""
    Reduces a dimension of a tensor by multiplying all elements in the dimension, by default. And also can
    reduce a dimension of `x` along the axis.

    Args:
        - **x** (Tensor[Number]) - The input tensor. The dtype of the tensor to be reduced is number.
          :math:`(N,*)` where :math:`*` means, any number of additional dimensions, its rank should be less than 8.
        - **axis** (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.
          Only constant value is allowed. Must be in the range [-rank(x), rank(x)).

    Outputs:
        Tensor, the dtype is the same as `x`.
    """
    if isinstance(axis, int):
        axis = (axis,)
    return torch.prod(x, dim=axis)

def concat_first_dim(input_x: Tensor) -> Tensor:
    r"""
    Concatenates a sequence of tensors along the first dimension.

    Args:
        input_x (Tensor): A sequence of tensors to be concatenated.

    Returns:
        Tensor, concatenated tensor.
    """
    return torch.cat(input_x, dim=0)

def concat_last_dim(input_x: Tensor) -> Tensor:
    r"""
    Concatenates a sequence of tensors along the last dimension.

    Args:
        input_x (Tensor): A sequence of tensors to be concatenated.

    Returns:
        Tensor, concatenated tensor.
    """
    return torch.cat(input_x, dim=-1)

def concat_penulti(input_x: Tensor) -> Tensor:
    r"""
    Concatenates a sequence of tensors along the penultimate dimension.

    Args:
        input_x (Tensor): A sequence of tensors to be concatenated.

    Returns:
        Tensor, concatenated tensor.
    """
    return torch.cat(input_x, dim=-2)

def stack_first_dim(input_x: Tensor) -> Tensor:
    r"""
    Stacks a sequence of tensors along the first dimension.

    Args:
        input_x (Tensor): A sequence of tensors to be stacked.

    Returns:
        Tensor, stacked tensor.
    """
    return torch.stack(input_x, dim=0)

def stack_last_dim(input_x: Tensor) -> Tensor:
    r"""
    Stacks a sequence of tensors along the last dimension.

    Args:
        input_x (Tensor): A sequence of tensors to be stacked.

    Returns:
        Tensor, stacked tensor.
    """
    return torch.stack(input_x, dim=-1)

def stack_penulti(input_x: Tensor) -> Tensor:
    r"""
    Stacks a sequence of tensors along the penultimate dimension.

    Args:
        input_x (Tensor): A sequence of tensors to be stacked.

    Returns:
        Tensor, stacked tensor.
    """
    return torch.stack(input_x, dim=-2)

def squeeze_first_dim(input_x: Tensor) -> Tensor:
    r"""
    Removes the first dimension of size 1 from the input tensor.

    Args:
        input_x (Tensor): The input tensor.

    Returns:
        Tensor, squeezed tensor.
    """
    return torch.squeeze(input_x, dim=0)

def squeeze_last_dim(input_x: Tensor) -> Tensor:
    r"""
    Removes the last dimension of size 1 from the input tensor.

    Args:
        input_x (Tensor): The input tensor.

    Returns:
        Tensor, squeezed tensor.
    """
    return torch.squeeze(input_x, dim=-1)

def squeeze_penulti(input_x: Tensor) -> Tensor:
    r"""
    Removes the penultimate dimension of size 1 from the input tensor.

    Args:
        input_x (Tensor): The input tensor.

    Returns:
        Tensor, squeezed tensor.
    """
    return torch.squeeze(input_x, dim=-2)

def periodic_variable(variable: Tensor,
                      upper: Tensor,
                      lower: Tensor = 0,
                      mask: Tensor = None,
                      ) -> Tensor:
    r"""
    Makes the variable periodic between lower and upper.

    Args:
        variable (Tensor): The input variable.
        upper (Tensor): The upper bound.
        lower (Tensor): The lower bound. Default: 0.
        mask (Tensor): The mask tensor. Default: None.

    Returns:
        Tensor, periodic variable.
    """
    period = upper - lower
    period_value = variable - period * torch.floor((variable - lower) / period)
    if mask is None:
        return period_value
    if mask.shape != variable.shape:
        mask = mask.expand_as(variable)  # 使用expand_as进行广播
    return torch.where(mask, period_value, variable)

def periodic_difference(difference: Tensor,
                        period: Tensor,
                        mask: Tensor = None,
                        offset: float = -0.5,
                        ) -> Tensor:
    r"""
    Makes the difference periodic.

    Args:
        difference (Tensor): The input difference.
        period (Tensor): The period.
        mask (Tensor): The mask tensor. Default: None.
        offset (float): The offset. Default: -0.5.

    Returns:
        Tensor, periodic difference.
    """

    period_diff = difference - period * torch.floor(difference / period - offset)
    if mask is None:
        return period_diff
    if mask.shape != difference.shape:
        mask = mask.expand_as(difference)  # 使用expand_as进行广播
    return torch.where(mask, period_diff, difference)

def gather_vector(tensor: Tensor, index: Tensor) -> Tensor:
    r"""Gather vector from the penultimate axis (`axis=-2`) of the tensor according to index.

    """
    # index = index.to(torch.int64)
    if index.shape[0] == 1:
        return torch.index_select(tensor, -2, index[0].view(-1)).reshape(tensor.shape[:-2]+index[0].shape[0:]+tensor.shape[-1:])
    if tensor.shape[0] == 1:
        return torch.index_select(tensor[0], -2, index.view(-1)).reshape(tensor[0].shape[:-2]+index.shape[0:]+tensor[0].shape[-1:])

    shape = index.shape+(tensor.shape[-1],)
    index = index.unsqueeze(-1).expand(*shape)
    tensor = tensor.unsqueeze(-2).expand(*shape)
    output = torch.gather(tensor, 1, index)

    return output

def gather_value(tensor: Tensor, index: Tensor) -> Tensor:
    r"""
    Gathers values from the input tensor according to the index.

    Args:
        tensor (Tensor): The input tensor.
        index (Tensor): The index tensor.

    Returns:
        Tensor, gathered values.
    """
    if index.shape[0] == 1:
        return torch.index_select(tensor, -1, index[0].view(-1)).reshape(tensor.shape[:-1]+index[0].shape[0:])
    if tensor.shape[0] == 1:
        return torch.index_select(tensor[0], -1, index.view(-1)).reshape(tensor[0].shape[:-1]+index.shape[0:])

    # # (B, N, M)
    # origin_shape = index.shape
    # # (B, N * M) <- (B, N, M)
    # index = torch.reshape(index, (origin_shape[0], -1))
    # # (B, N * M)
    # values = torch.gather(tensor, -1, index)
    # (B, N, M)
    # return torch.reshape(values, origin_shape)
    # output = torch.index_select(tensor, -1, index.view(-1)).reshape(tensor.shape[:-1]+index.shape[0:])
    shape = tensor.shape + (index.shape[-1],)
    output = torch.gather(tensor.unsqueeze(-1).expand(*shape), 1, index)
    return output

def pbc_box_reshape(pbc_box: Tensor, ndim: int) -> Tensor:
    r"""
    Reshapes the periodic boundary condition box.

    Args:
        pbc_box (Tensor): Periodic boundary condition box.
        ndim (int): Number of dimensions.

    Returns:
        Tensor, reshaped periodic boundary condition box.
    """
    if pbc_box is None:
        # 返回一个全1的tensor，维度与position一致
        return torch.ones(ndim)
        
    shape = pbc_box.shape
    if len(shape) == 1:
        pbc_box = pbc_box.unsqueeze(0)
    shape = pbc_box.shape
    if len(shape) + 1 == ndim:
        pbc_box = pbc_box.unsqueeze(1)
    elif len(shape) + 2 == ndim:
        pbc_box = pbc_box.unsqueeze(1).unsqueeze(1)
    return pbc_box

def pbc_image(position: Tensor, pbc_box: Tensor, offset: float = 0) -> Tensor:
    r"""
    Calculates the periodic boundary condition image.

    Args:
        position (Tensor): Position coordinates.
        pbc_box (Tensor): Periodic boundary condition box.
        offset (float): Offset value. Default: 0

    Returns:
        Tensor, periodic boundary condition image.
    """
    if pbc_box is None:
        # 如果没有周期性边界条件，返回零图像
        return torch.zeros_like(position)
        
    pbc_box = pbc_box_reshape(pbc_box, position.ndim)
    return -torch.floor(position / pbc_box - offset)

def coordinate_in_pbc(position: Tensor, pbc_box: Tensor, offset: float = 0) -> Tensor:
    r"""
    Calculates the coordinate in periodic boundary condition.

    Args:
        position (Tensor): The position tensor.
        pbc_box (Tensor): The periodic boundary condition box.
        offset (float): The offset. Default: 0.

    Returns:
        Tensor, coordinate in periodic boundary condition.
    """
    pbc_box = pbc_box_reshape(pbc_box, position.ndim)
    return position - pbc_box * torch.floor(position / pbc_box - offset)

def vector_in_pbc(vector: Tensor, pbc_box: Tensor, offset: float = -0.5) -> Tensor:
    r"""
    Calculates the vector in periodic boundary condition.

    Args:
        vector (Tensor): The vector tensor.
        pbc_box (Tensor): The periodic boundary condition box.
        offset (float): The offset. Default: -0.5.

    Returns:
        Tensor, vector in periodic boundary condition.
    """
    pbc_box = pbc_box_reshape(pbc_box, vector.ndim)
    box_nograd = pbc_box.detach()
    inv_box = torch.reciprocal(box_nograd)
    vector = vector - box_nograd * torch.floor(vector * inv_box - offset)
    return vector * inv_box * pbc_box

def calc_vector_nopbc(initial: Tensor, terminal: Tensor) -> Tensor:
    r"""
    Calculates the vector without periodic boundary condition.

    Args:
        initial (Tensor): The initial position tensor.
        terminal (Tensor): The terminal position tensor.

    Returns:
        Tensor, vector without periodic boundary condition.
    """
    return terminal - initial

def calc_vector_pbc(initial: Tensor, terminal: Tensor, pbc_box: Tensor) -> Tensor:
    r"""
    Calculates the vector with periodic boundary condition.

    Args:
        initial (Tensor): The initial position tensor.
        terminal (Tensor): The terminal position tensor.
        pbc_box (Tensor): The periodic boundary condition box.

    Returns:
        Tensor, vector with periodic boundary condition.
    """
    vector = calc_vector_nopbc(initial, terminal)
    return vector_in_pbc(vector, pbc_box)

def calc_vector(initial: Tensor, terminal: Tensor, pbc_box: Tensor = None) -> Tensor:
    r"""
    Calculates the vector with or without periodic boundary condition.

    Args:
        initial (Tensor): The initial position tensor.
        terminal (Tensor): The terminal position tensor.
        pbc_box (Tensor): The periodic boundary condition box. Default: None.

    Returns:
        Tensor, vector with or without periodic boundary condition.
    """
    if pbc_box is None:
        return calc_vector_nopbc(initial, terminal)
    return calc_vector_pbc(initial, terminal, pbc_box)

def calc_distance_nopbc(position_a: Tensor,
                        position_b: Tensor,
                        keepdims: bool = False,
                        ) -> Tensor:
    r"""
    Calculates the distance without periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, distance without periodic boundary condition.
    """
    vector = calc_vector_nopbc(position_a, position_b)
    return torch.norm(vector, dim=-1, keepdim=keepdims)

def calc_distance_pbc(position_a: Tensor,
                      position_b: Tensor,
                      pbc_box: Tensor = None,
                      keepdims: bool = False
                      ) -> Tensor:
    r"""
    Calculates the distance with periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        pbc_box (Tensor): The periodic boundary condition box. Default: None.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, distance with periodic boundary condition.
    """
    vector = calc_vector_pbc(position_a, position_b, pbc_box)
    return torch.norm(vector, dim=-1, keepdim=keepdims)

def calc_distance(position_a: Tensor,
                  position_b: Tensor,
                  pbc_box: Tensor = None,
                  keepdims: bool = False,
                  ) -> Tensor:
    r"""
    Calculates the distance with or without periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        pbc_box (Tensor): The periodic boundary condition box. Default: None.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, distance with or without periodic boundary condition.
    """
    if pbc_box is None:
        return calc_distance_nopbc(position_a, position_b, keepdims)
    return calc_distance_pbc(position_a, position_b, pbc_box, keepdims)

def calc_angle_by_vectors(vector1: Tensor,
                          vector2: Tensor,
                          keepdims: bool = False
                          ) -> Tensor:
    r"""
    Calculates the angle between two vectors.

    Args:
        vector1 (Tensor): The first vector.
        vector2 (Tensor): The second vector.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, angle between two vectors.
    """
    dot = torch.sum(vector1 * vector2, dim=-1, keepdim=keepdims)
    norm1 = torch.norm(vector1, dim=-1, keepdim=keepdims)
    norm2 = torch.norm(vector2, dim=-1, keepdim=keepdims)
    return torch.acos(dot / (norm1 * norm2))

def calc_angle_nopbc(position_a: Tensor,
                     position_b: Tensor,
                     position_c: Tensor,
                     keepdims: bool = False,
                     ) -> Tensor:
    r"""
    Calculates the angle without periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        position_c (Tensor): The position tensor c.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, angle without periodic boundary condition.
    """
    vector1 = calc_vector_nopbc(position_b, position_a)
    vector2 = calc_vector_nopbc(position_b, position_c)
    return calc_angle_by_vectors(vector1, vector2, keepdims)

def calc_angle_pbc(position_a: Tensor,
                   position_b: Tensor,
                   position_c: Tensor,
                   pbc_box: Tensor,
                   keepdims: bool = False,
                   ) -> Tensor:
    r"""
    Calculates the angle with periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        position_c (Tensor): The position tensor c.
        pbc_box (Tensor): The periodic boundary condition box.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, angle with periodic boundary condition.
    """
    vector1 = calc_vector_pbc(position_b, position_a, pbc_box)
    vector2 = calc_vector_pbc(position_b, position_c, pbc_box)
    return calc_angle_by_vectors(vector1, vector2, keepdims)

def calc_angle(position_a: Tensor,
               position_b: Tensor,
               position_c: Tensor,
               pbc_box: Tensor = None,
               keepdims: bool = False,
               ) -> Tensor:
    r"""
    Calculates the angle with or without periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        position_c (Tensor): The position tensor c.
        pbc_box (Tensor): The periodic boundary condition box. Default: None.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, angle with or without periodic boundary condition.
    """
    if pbc_box is None:
        return calc_angle_nopbc(position_a, position_b, position_c, keepdims)
    return calc_angle_pbc(position_a, position_b, position_c, pbc_box, keepdims)

def calc_torsion_by_vectors(vector1: Tensor,
                            vector2: Tensor,
                            axis_vector: Tensor = None,
                            keepdims: bool = False,
                            ) -> Tensor:
    r"""
    Calculates the torsion angle between two vectors.

    Args:
        vector1 (Tensor): The first vector.
        vector2 (Tensor): The second vector.
        axis_vector (Tensor): The axis vector. Default: None.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, torsion angle between two vectors.
    """
    if axis_vector is None:
        axis_vector = torch.cross(vector1, vector2, dim=-1)
    normal1 = torch.cross(vector1, axis_vector, dim=-1)
    normal2 = torch.cross(vector2, axis_vector, dim=-1)
    dot = torch.sum(normal1 * normal2, dim=-1, keepdim=keepdims)
    norm1 = torch.norm(normal1, dim=-1, keepdim=keepdims)
    norm2 = torch.norm(normal2, dim=-1, keepdim=keepdims)
    return torch.acos(dot / (norm1 * norm2))

def calc_torsion_nopbc(position_a: Tensor,
                       position_b: Tensor,
                       position_c: Tensor,
                       position_d: Tensor,
                       keepdims: bool = False,
                       ) -> Tensor:
    r"""
    Calculates the torsion angle without periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        position_c (Tensor): The position tensor c.
        position_d (Tensor): The position tensor d.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, torsion angle without periodic boundary condition.
    """
    vector1 = calc_vector_nopbc(position_b, position_a)
    vector2 = calc_vector_nopbc(position_c, position_d)
    axis_vector = calc_vector_nopbc(position_b, position_c)
    return calc_torsion_by_vectors(vector1, vector2, axis_vector, keepdims)

def calc_torsion_pbc(position_a: Tensor,
                     position_b: Tensor,
                     position_c: Tensor,
                     position_d: Tensor,
                     pbc_box: Tensor,
                     keepdims: bool = False,
                     ) -> Tensor:
    r"""
    Calculates the torsion angle with periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        position_c (Tensor): The position tensor c.
        position_d (Tensor): The position tensor d.
        pbc_box (Tensor): The periodic boundary condition box.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, torsion angle with periodic boundary condition.
    """
    vector1 = calc_vector_pbc(position_b, position_a, pbc_box)
    vector2 = calc_vector_pbc(position_c, position_d, pbc_box)
    axis_vector = calc_vector_pbc(position_b, position_c, pbc_box)
    return calc_torsion_by_vectors(vector1, vector2, axis_vector, keepdims)

def calc_torsion(position_a: Tensor,
                 position_b: Tensor,
                 position_c: Tensor,
                 position_d: Tensor,
                 pbc_box: Tensor = None,
                 keepdims: bool = False,
                 ) -> Tensor:
    r"""
    Calculates the torsion angle with or without periodic boundary condition.

    Args:
        position_a (Tensor): The position tensor a.
        position_b (Tensor): The position tensor b.
        position_c (Tensor): The position tensor c.
        position_d (Tensor): The position tensor d.
        pbc_box (Tensor): The periodic boundary condition box. Default: None.
        keepdims (bool): Whether to keep the last dimension. Default: False.

    Returns:
        Tensor, torsion angle with or without periodic boundary condition.
    """
    if pbc_box is None:
        return calc_torsion_nopbc(position_a, position_b, position_c, position_d, keepdims)
    return calc_torsion_pbc(position_a, position_b, position_c, position_d, pbc_box, keepdims)

def coulomb_interaction(q_i: Tensor,
                        q_j: Tensor,
                        r_ij: Tensor,
                        mask: Tensor = None,
                        coulomb_const: float = 1,
                        ):
    r"""
    Calculates the Coulomb interaction.

    Args:
        q_i (Tensor): The charge tensor i.
        q_j (Tensor): The charge tensor j.
        r_ij (Tensor): The distance tensor.
        mask (Tensor): The mask tensor. Default: None.
        coulomb_const (float): The Coulomb constant. Default: 1.

    Returns:
        Tensor, Coulomb interaction.
    """
    energy = coulomb_const * q_i * q_j / r_ij

    if mask is None:
        return energy
    return energy * mask

def lennard_jones_potential(epsilon: Tensor, sigma: Tensor, r_ij: Tensor, mask: Tensor = None) -> Tensor:
    r"""
    Calculates the Lennard-Jones potential.

    Args:
        epsilon (Tensor): The epsilon tensor.
        sigma (Tensor): The sigma tensor.
        r_ij (Tensor): The distance tensor.
        mask (Tensor): The mask tensor. Default: None.

    Returns:
        Tensor, Lennard-Jones potential.
    """
    r6 = (sigma / r_ij) ** 6
    energy = 4 * epsilon * (r6 ** 2 - r6)

    if mask is None:
        return energy
    return energy * mask

def lennard_jones_potential2(epsilon: Tensor, r_0: Tensor, r_ij: Tensor, mask: Tensor = None) -> Tensor:
    r"""
    Calculates the Lennard-Jones potential with r_0.

    Args:
        epsilon (Tensor): The epsilon tensor.
        r_0 (Tensor): The r_0 tensor.
        r_ij (Tensor): The distance tensor.
        mask (Tensor): The mask tensor. Default: None.

    Returns:
        Tensor, Lennard-Jones potential with r_0.
    """
    
    r6 = (r_0 / r_ij) ** 6
    energy = epsilon * (r6 ** 2 - 2 * r6)

    if mask is None:
        return energy
    return energy * mask

def get_integer(value: Union[int, Tensor, Parameter, ndarray, None]) -> int:
    r"""
    Gets the integer value.

    Args:
        value (Union[int, Tensor, Parameter, ndarray, None]): The input value.

    Returns:
        int, integer value.
    """
    if value is None:
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, (Tensor, Parameter)):
        return int(value.item())
    if isinstance(value, ndarray):
        return int(value.item())
    raise TypeError(f'Unsupported type: {type(value)}')

def get_ndarray(value: Union[Tensor, Parameter, ndarray, List[float], Tuple[float]],
                dtype: type = None) -> ndarray:
    r"""
    Gets the numpy array.

    Args:
        value (Union[Tensor, Parameter, ndarray, List[float], Tuple[float]]): The input value.
        dtype (type): The data type. Default: None.

    Returns:
        ndarray, numpy array.
    """
    if isinstance(value, (Tensor, Parameter)):
        value = value.detach().cpu().numpy()
    if isinstance(value, (list, tuple)):
        value = np.array(value)
    if dtype is not None:
        value = value.astype(dtype)
    return value

def get_tensor(value, dtype=None):
    """Convert value to tensor."""
    if value is None:
        return None
    if isinstance(value, torch.Tensor):
        if dtype is not None:
            value = value.to(dtype)
        return value.to(get_global_device())
    if isinstance(value, np.ndarray):
        value = torch.from_numpy(value)
    else:
        try:
            value = torch.tensor(value)
        except (TypeError, RuntimeError):
            raise ValueError(f"Cannot convert {type(value)} to tensor")
    if dtype is not None:
        value = value.to(dtype)
    return value.to(get_global_device())

def check_broadcast(shape0: tuple, shape1: tuple) -> tuple:
    r"""
    Checks if two shapes can be broadcast.

    Args:
        shape0 (tuple): The first shape.
        shape1 (tuple): The second shape.

    Returns:
        tuple, broadcast shape.
    """
    if shape0 is None:
        return shape1
    if shape1 is None:
        return shape0
    
    if len(shape0) < len(shape1):
        shape0 = (1,) * (len(shape1) - len(shape0)) + shape0
    if len(shape1) < len(shape0):
        shape1 = (1,) * (len(shape0) - len(shape1)) + shape1
    shape = []
    for s0, s1 in zip(shape0, shape1):
        if s0 == 1:
            shape.append(s1)
        elif s1 == 1:
            shape.append(s0)
        elif s0 == s1:
            shape.append(s0)
        else:
            raise ValueError(f'Cannot broadcast shapes {shape0} and {shape1}')
    return tuple(shape)

def any_none(iterable: Iterable) -> bool:
    r"""
    Checks if any element is None.

    Args:
        iterable (Iterable): The input iterable.

    Returns:
        bool, whether any element is None.
    """
    return any(x is None for x in iterable)

def all_none(iterable: Iterable) -> bool:
    r"""
    Checks if all elements are None.

    Args:
        iterable (Iterable): The input iterable.

    Returns:
        bool, whether all elements are None.
    """
    return all(x is None for x in iterable)

def any_not_none(iterable: Iterable) -> bool:
    r"""
    Checks if any element is not None.

    Args:
        iterable (Iterable): The input iterable.

    Returns:
        bool, whether any element is not None.
    """
    return any(x is not None for x in iterable)

def all_not_none(iterable: Iterable) -> bool:
    r"""
    Checks if all elements are not None.

    Args:
        iterable (Iterable): The input iterable.

    Returns:
        bool, whether all elements are not None.
    """
    return all(x is not None for x in iterable)

def get_arguments(locals_: dict, kwargs: dict = None) -> dict:
    r"""
    Gets the arguments.

    Args:
        locals_ (dict): The locals dictionary.
        kwargs (dict): The kwargs dictionary. Default: None.

    Returns:
        dict, arguments.
    """
    args = {}
    for key, value in locals_.items():
        if key != 'kwargs' and key != 'self' and key != '__class__':
            args[key] = value
    if kwargs is not None:
        args.update(kwargs)
    return args

def _bonds_in(bonds, bond):
    """ Check bonds exists in both sets. """
    return (bonds == bond).all(-1)


def bonds_in(bonds, batch_bond):
    """ Return if batch_bond exists in bonds.
    Args:
        bonds: The total bonds set.
        batch_bond: The input bond set.

    Returns:
        If batch_bond exists in bonds, the mask will be 1, else 0.
    """
    return torch.vmap(_bonds_in, in_dims=(None, -2), out_dims=0)(bonds, batch_bond).sum(dim=-3)

def calc_gaussian(x, mu=0):
    r"""
    Calculates the Gaussian function.

    Args:
        x (Tensor): The input tensor.
        mu (float): The mean. Default: 0.

    Returns:
        Tensor, Gaussian function.
    """
    return torch.exp(-0.5 * (x - mu) ** 2)

class H5MDProcessor:
    def __init__(self, path):
        self.path = path
        self.time = None
        self.bias = None
        self.cv = None
        self.coord = None
        self.pbc_box = None
        self.atom_name = None
        self.atomic_number = None
        self.temperature = None
        self.grids0 = None
        self.hills0 = None

    def _process_single_bead(self, traj_key):
        """
        处理每个 bead 的数据，提取所需信息。
        """
        with h5py.File(self.path) as t:
            data = t[f'particles/{traj_key}/position/value'][:][:,None]
            ts = t[f'particles/{traj_key}/position/time'][:].reshape(-1,1)
            tp = t[f'observables/{traj_key}/temperature/value'][:][:,None]
             # 检查并提取bias和cv
            box = t[f'particles/{traj_key}/box/edges/value'][:][:,None] if 'edges' in t[f'particles/{traj_key}/box'] else None
            bs = t[f'observables/{traj_key}/bias_potential/value'][:][:,None] if 'bias_potential' in t[f'observables/{traj_key}'] else None
            cv = t[f'/observables/{traj_key}/cv/value'][:][:,None] if 'cv' in t[f'observables/{traj_key}'] else None

        return data, ts, tp, box, bs, cv

    def process(self, bead='all'):
        """
        处理指定 bead 的数据
        """
        with h5py.File(self.path) as t:
            trajectory_keys = list(t['particles'].keys())

            if isinstance(bead, int):
                trajectory_keys = [f'trajectory{bead}']
            elif isinstance(bead, list):
                trajectory_keys = [f'trajectory{b}' for b in bead]
            elif bead == 'all':
                pass  # 不进行过滤，直接处理所有trajectory

            # 并行处理所有选定的trajectory
            results = Parallel(n_jobs=-1)(delayed(self._process_single_bead)(key) for key in trajectory_keys)
            # 从并行结果中提取数据
            coord, time, temperature, box, bias, cv = zip(*results)
            self.atom_name = np.array(t['parameters']['vmd_structure']['name'], dtype=str)
            self.atomic_number = np.array(t['parameters']['vmd_structure']['atomicnumber'], dtype=int)

            if 'meta_potential' in t['parameters']:
                try:
                    self.grids0 = np.array(t['parameters/meta_potential/grids'], dtype=np.float32)
                    self.hills0 = np.array(t['parameters/meta_potential/hills'], dtype=np.float32)
                except:
                    self.grids0 = np.array(t['parameters/meta_potential/grids0'], dtype=np.float32)
                    self.hills0 = np.array(t['parameters/meta_potential/hills0'], dtype=np.float32)

        # 合并所有trajectory的数据
        self.time = np.concatenate(time)
        self.temperature = np.concatenate(temperature)
        self.coord = np.concatenate(coord, axis=1)
        # 检查并合并bias和cv
        if bias[0] is not None:
            self.bias = np.concatenate(bias, axis=1)
        if cv[0] is not None:
            self.cv = np.concatenate(cv, axis=1)
        if box[0] is not None:
            self.pbc_box = np.concatenate(box, axis=1)
        return self
    
    def toxyz(self, frame, traj_id=0):
            # 统一 atom 维度
        atom = self.atom_name
        if atom.ndim == 2 and atom.shape[0] == 1:
            atom = atom[0]
        elif atom.ndim > 2:
            raise ValueError("atom should be of shape (Natom), (Nframe, Natom) or (1, Natom)")

        if isinstance(frame, str) and frame == 'all':
            data = self.coord[:]
        else:
            data = self.coord[frame]
            
        if traj_id != 'all':
            data = data[:, traj_id]

        if data.ndim == 2:
            data = data[np.newaxis, ...]  # (Natom, Ndim) 转换为 (1, Natom, Ndim)
        elif data.ndim != 3:
            raise ValueError("data should be of shape (Nframe, Natom, Ndim) or (Natom, Ndim)")

        num_frames, num_atoms, ndim = data.shape

        if atom.ndim == 1 and atom.size != num_atoms:
            raise ValueError("Length of atom array does not match number of atoms in data")
        elif atom.ndim == 2 and atom.shape != (num_frames, num_atoms):
            raise ValueError("Shape of atom array does not match data shape")

        # 处理 box 参数
        box_str = ""
        pbc_str = ""

        def process_box(box):
            box = np.asarray(box).reshape(-1)
            if len(box) == 1:
                box_str = f'Lattice="{box[0]} 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" '
            elif len(box) == 2:
                box_str = f'Lattice="{box[0]} 0.0 0.0 0.0 {box[1]} 0.0 0.0 0.0 0.0" '
            elif len(box) == 3:
                box_str = f'Lattice="{box[0]} 0.0 0.0 0.0 {box[1]} 0.0 0.0 0.0 {box[2]}" '
            else:
                raise ValueError("box should be of shape (Ndim) where Ndim is 1, 2 or 3")
            pbc_str = "pbc=" + "T "*len(box)
            return box_str, pbc_str

        xyz_contents = []
        for i in range(num_frames):
            current_atom = atom if atom.ndim == 1 else atom[i]
            frame_data = data[i] * 10
            # 补全不足 3 维的坐标
            padded_data = np.pad(frame_data, ((0, 0), (0, 3 - ndim)), 'constant', constant_values=0)

            properties_str = f"Properties=species:S:1:pos:R:{ndim}"
            box_str, pbc_str = process_box(self.pbc_box[i] * 10) if self.pbc_box is not None else ("", "")
            content = f"{num_atoms}\n{box_str}{properties_str} {pbc_str}\n"
            atom_block = "\n".join(
                f"{a}     {x:.8f}     {y:.8f}     {z:.8f}"
                for a, (x, y, z) in zip(current_atom, padded_data)
            )
            xyz_contents.append(content + atom_block)

        return "\n".join(xyz_contents)
    
class TorchSPONGEDevice:
    def __init__(self) -> None:
        self._device = torch.device('cpu')

    def set_device(self, device):
        self._device = device

    @property
    def device(self):
        """获取当前全局设备"""
        return self._device

_GLOBAL_DEVICE = TorchSPONGEDevice()

def get_global_device():
    """获取全局设备"""
    return _GLOBAL_DEVICE.device

def set_torchsponge_global_device(device):
    _GLOBAL_DEVICE.set_device(device)
    print(f'Set TorchSPONGE global device to {get_global_device()}')
